#include "raytracer.h"
#include "material.h"
#include "vectors.h"
#include "argparser.h"
#include "raytree.h"
#include "utils.h"
#include "mesh.h"
#include "face.h"
#include "primitive.h"

// ===========================================================================
// casts a single ray through the scene geometry and finds the closest hit
bool RayTracer::CastRay(Ray &ray, Hit &h, double time, bool use_rasterized_patches) const {
	bool answer = false;

	// intersect each of the true quads 
	for (int i = 0; i < mesh->numOriginalQuads(); i++)
	{
		Face *f = mesh->getOriginalQuad(i);

		if (f->intersect(ray,h,args->intersect_backfacing)) answer = true;
	}

	// intersect each of the primitives (either the patches, or the original primitives)
	int num_primitives = mesh->numPrimitives();

	for (int i = 0; i < num_primitives; i++)
	{
		if (mesh->getPrimitive(i)->intersect(ray, h, time)) answer = true;
	}

	return answer;
}

// ===========================================================================
// does the recursive (shadow rays & recursive rays) work
Vec3f RayTracer::TraceRay(Ray &ray, Hit &hit, double time, int bounce_count, bool glossyRay) const
{
	hit = Hit();

	// First cast a ray and see if we hit anything.
	bool intersect = CastRay(ray, hit, time, false);
  
	// if there is no intersection, simply return the background color
	if (!intersect)
	{
		const Vec3f &bg = mesh->getBackgroundColor();

		return Vec3f(srgb_to_linear(bg.r()), srgb_to_linear(bg.g()), srgb_to_linear(bg.b()));    
	}

	// otherwise decide what to do based on the material
	Material *m = hit.getMaterial();
	assert (m != NULL);

	// rays coming from the light source are set to white, don't bother to ray trace further.
	if (m->getEmittedColor().Length() > 0.001)
	{
		return Vec3f(1,1,1);
	} 
  
	Vec3f normal = hit.getNormal();
	Vec3f point = ray.pointAtParameter(hit.getT());

	// ----------------------------------------------
	// start with the indirect light (ambient light)
	Vec3f diffuse_color = m->getDiffuseColor(hit.get_s(),hit.get_t());

	// the usual ray tracing hack for indirect light
	Vec3f answer = diffuse_color * args->ambient_light;

	// add contributions from each light that is not in shadow
	int num_lights = mesh->getLights().size();

	for (int i = 0; i < num_lights; i++)
	{
		Face *f = mesh->getLights()[i];
		Vec3f pointOnLight = f->computeCentroid();
		Vec3f lightColor = f->getMaterial()->getEmittedColor() * f->getArea();
		Vec3f dirToLight = pointOnLight - point;
		double dist = dirToLight.Length();
		dirToLight.Normalize();

		lightColor /= (M_PI * dist * dist);

		double shadowIntersects = 0;
		int shadowSamples = 1;

		if (!glossyRay)
		{
			shadowSamples = args->num_shadow_samples;
		}

		// Cast n samples for each shadow
		for (int s=0; s < shadowSamples; s++)
		{
			Vec3f toLight = f->RandomPoint() - point;
			dist = toLight.Length();
			toLight.Normalize();

			// Compute the new shadow ray
			Ray shadowRay(point + toLight * 0.001, toLight);
			Hit shadowHit = Hit();

			// Cast the shadow ray
			CastRay(shadowRay, shadowHit, time, false);
			RayTree::AddShadowSegment(shadowRay, 0.0, dist);

			if (shadowHit.getT() < 0.99 * dist)
			{
				shadowIntersects++;
			}
			
			if (s == args->num_shadow_cutoff && shadowIntersects == 0)
			{
				break;
			}
		}

		double shadowPercent = (1.0 - shadowIntersects /shadowSamples);

		// Only shade non-shadows
		answer += m->Shade(ray, hit, dirToLight, lightColor * shadowPercent, args);
	}

	// Add refraction contribution
	if (m->getIOR() != 1.0 && bounce_count == 0)
	{
		Vec3f refractionDirection = RefractionDirection(normal, ray.getDirection(), m->getIOR());
		Ray refractRay(point + refractionDirection * 0.001, refractionDirection);

		Hit refractHit;
		Vec3f refractColor = TraceRay(refractRay, refractHit, time, 1, false); 

		RayTree::AddReflectedSegment(refractRay, 0.0, refractHit.getT());

		Vec3f rNormal = -refractHit.getNormal();
		Vec3f rPoint = refractRay.pointAtParameter(refractHit.getT());

		refractionDirection = RefractionDirection(rNormal, refractRay.getDirection(), m->getIOR());

		Ray outputRay(rPoint + refractionDirection * 0.001, refractionDirection);

		Vec3f outputColor = TraceRay(outputRay, refractHit, time, 1, false); 

		RayTree::AddReflectedSegment(outputRay, 0.0, refractHit.getT());

		answer += outputColor * 0.5;
	}
      
	// add contribution from reflection, if the surface is shiny
	Vec3f reflectiveColor = m->getReflectiveColor();

	if (reflectiveColor == Vec3f(0,0,0)) { return answer; }

	if (bounce_count < args->num_bounces)
	{
		Vec3f mirrorReflection = MirrorDirection(normal, ray.getDirection());
		Vec3f sphereCenter;

		Hit reflectHit;

		double glossy = m->getGlossy();
		int reflect_iters = 1;
		double glossy_factor = 1.0;

		if (glossy > 0 && !glossyRay)
		{
			double sphereDistance = 1.0 / glossy;
			reflect_iters = 100 * exp(-sphereDistance / 4.0) + 1;
			glossy_factor = 1.0 / reflect_iters;

			sphereCenter = point + mirrorReflection * sphereDistance; 
		}

		for (int i=0; i < reflect_iters; i++)
		{
			Vec3f reflectColor;

			if (glossy == 0)
			{
				Ray reflectRay(point + mirrorReflection * 0.001, mirrorReflection);

				reflectColor = TraceRay(reflectRay, reflectHit, time, bounce_count + 1, false);

				RayTree::AddReflectedSegment(reflectRay, 0.0, reflectHit.getT());
			}
			else
			{
				Vec3f pointOnSphere = sphereCenter + RandomUnitVector();
				Vec3f reflectDirection = pointOnSphere - point;

				reflectDirection.Normalize();

				Ray reflectRay(point + reflectDirection * 0.001, reflectDirection);
				reflectColor = TraceRay(reflectRay, reflectHit, time, bounce_count + 1, true) * glossy_factor;

				RayTree::AddReflectedSegment(reflectRay, 0.0, reflectHit.getT());
			}

			answer += Vec3f(reflectColor.x() * reflectiveColor.x(),
							reflectColor.y() * reflectiveColor.y(),
							reflectColor.z() * reflectiveColor.z());
		}
	}
  
	return answer; 
}
